package com.example.fileupload;

import com.example.Constant;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * @author zhanghui
 * @date 2021/12/6 13:59
 */
public class FileUploadServer {
    public static void main(String[] args) throws IOException {
        FileUploadServer server = new FileUploadServer();
        server.startServer();
    }

    static class Client {
        String fileName;
        long fileLength;
        // 开始传输时间
        long startTime;
        // 客户端地址
        InetSocketAddress remoteAddress;
        // 文件输出通道
        FileChannel outChannel;
    }

    private Charset charset = Charset.forName("UTF-8");
    private ByteBuffer buffer = ByteBuffer.allocate(1024);
    Map<SelectableChannel, Client> clientMap = new HashMap<>();

    public void startServer() throws IOException {
        // 1、获取Selector选择器
        Selector selector = Selector.open();

        // 2、获取通道
        ServerSocketChannel serverChannel = ServerSocketChannel.open();
        ServerSocket serverSocket = serverChannel.socket();

        // 3.设置为非阻塞
        serverChannel.configureBlocking(false);
        // 4、绑定连接
        InetSocketAddress address = new InetSocketAddress(Constant.SERVER_PORT);
        serverSocket.bind(address);
        // 5、将通道注册到选择器上,并注册的IO事件为：“接收新连接”
        serverChannel.register(selector, SelectionKey.OP_ACCEPT);
        System.out.println("服务端已启动");

        while (selector.select() > 0) {
            Iterator<SelectionKey> selectKeys = selector.selectedKeys().iterator();
            if (selectKeys.hasNext()) {
                SelectionKey selectionKey = selectKeys.next();

                if (selectionKey.isAcceptable()) {
                    // 若是连接事件, 则获取通道并记录客户端的信息
                    ServerSocketChannel serverSocketChannel = (ServerSocketChannel) selectionKey.channel();
                    SocketChannel socketChannel = serverSocketChannel.accept();
                    if (socketChannel == null) {
                        continue;
                    }
                    socketChannel.configureBlocking(false);
                    socketChannel.register(selector, SelectionKey.OP_READ);
                    // 记录客户端信息
                    Client client = new Client();

                    client.remoteAddress = (InetSocketAddress) socketChannel.getRemoteAddress();
                    clientMap.put(socketChannel, client);
                    System.out.println(client.remoteAddress + " 连接成功");
                } else if (selectionKey.isReadable()) {
                    receiveFile(selectionKey);
                }

                selectKeys.remove();
            }
        }
    }

    /**
     * 业务处理, 从通道读取数据并写入文件
     *
     * @param selectionKey
     */
    private void receiveFile(SelectionKey selectionKey) throws IOException {
        Client client = clientMap.get(selectionKey.channel());

        SocketChannel socketChannel = (SocketChannel) selectionKey.channel();
        int readLength = 0;
        buffer.clear();
        while ((readLength = socketChannel.read(buffer)) > 0) {
            buffer.flip();
            // 客户端发送的以为文件名, 文件长度, 文件数据
            if (client.fileName == null) {
                String fileName = charset.decode(buffer).toString();
                File file = new File(fileName);
                System.out.println("接收文件[" + fileName + "]");
                FileOutputStream fileOutputStream = new FileOutputStream(file);
                client.outChannel = fileOutputStream.getChannel();
            } else if (client.fileLength == 0) {
                client.fileLength = buffer.getLong();
                client.startTime = System.currentTimeMillis();
                System.out.println("开始接收[" + client.fileName + "]");
            } else {
                client.outChannel.write(buffer);
            }
            buffer.clear();
        }
        selectionKey.cancel();

        if (readLength == -1) {
            long end = System.currentTimeMillis();
            System.out.println("[" + client.fileName + "]保存完成, 大小: " + client.fileLength + ", 用时: " + (end - client.startTime) * 1.0 / 1000 + "秒");
        }
    }
}
